!pip install pyro-pplBayesian LSTM for volatility forecasting
Most market forecasts hand you a single number. This one tries to hand you a number and an honest sense of how much it could wiggle. In finance, knowing you’re uncertain is often more useful than being precisely wrong, so the project pairs a deterministic LSTM with a Bayesian output layer. The LSTM learns the sequence of market conditions; the Bayesian layer turns the final prediction into a distribution rather than a point. Pyro handles the variational inference, and Optuna hunts for the settings that actually perform.
!pip install optunaGetting the data ready
I reach for the S&P 500 index and pull its full history. On top of the close price I engineer the features the model will care about: log returns, rolling volatility, short and long moving averages, RSI, and percentage changes in volume. Everything is standardized, then cut into 21-day windows so the LSTM has a clear sequence to chew on.
What the model looks like
At the front is a plain LSTM that reads the windows and tracks temporal structure. The twist comes after: the final linear layer is Bayesian, with priors on its weights, bias, and noise term. Rather than learning one fixed set of weights, the model learns a distribution over them, and that distribution is exactly what turns a single prediction into a forecast with a spread.
Finding the right settings
Deep networks are finicky, so I let Optuna tune the hidden dimension and the learning rate. Each candidate configuration trains on an inner split and gets scored on validation loss, with early stopping to keep things from dragging on. The result is a model tuned on evidence rather than guesswork.
Training
Once the hyperparameters are locked in, the final model is retrained on the full training set. The ELBO loss — Pyro’s objective for variational inference — gives a running read on how the fit is settling, and it’s worth watching it come down.
Checking the data’s behavior
Before trusting any forecast, a couple of sanity checks are worth running: how much does volatility persist, and is the series stationary enough to model? An augmented Dickey–Fuller test answers that last one.
Forecasting volatility
With the model trained, I sample from the posterior to build an out-of-sample forecast. Each sample is one plausible future; averaging them gives a mean prediction, and the spread around that mean is the model’s stated uncertainty. Mean squared error and mean absolute error are reported against the held-out volatility.
import torch, numpy as np, pandas as pd, yfinance as yf, pyro, warnings, optuna
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset
from pyro.nn import PyroModule, PyroSample
import pyro.distributions as dist
from pyro.infer import SVI, Trace_ELBO, Predictive
from pyro.optim import Adam
from pyro.infer.autoguide import AutoNormal
from sklearn.preprocessing import StandardScaler
from statsmodels.tsa.stattools import adfuller
from sklearn.metrics import mean_squared_error, mean_absolute_error
warnings.filterwarnings("ignore")
torch.manual_seed(42); pyro.set_rng_seed(42)
# -------------------------
# Data Retrieval & Feature Engineering
# -------------------------
ticker = "^GSPC"
df = yf.Ticker(ticker).history(period="max", auto_adjust=True).reset_index()
if df.empty: raise ValueError("No data retrieved")
df['Log_Return'] = np.log(df['Close']+1e-9).diff()
df['Volatility'] = df['Log_Return'].rolling(21, min_periods=1).std()
df['MA_10'] = df['Close'].rolling(10).mean()
df['MA_50'] = df['Close'].rolling(50).mean()
delta = df['Close'].diff()
gain = delta.clip(lower=0)
loss = -delta.clip(upper=0)
avg_gain = gain.rolling(14, min_periods=1).mean()
avg_loss = loss.rolling(14, min_periods=1).mean()
df['RSI'] = 100 - 100/(1+avg_gain/(avg_loss+1e-8))
df['Volume_Change'] = df['Volume'].replace(0,np.nan).pct_change().fillna(0)
cols = ['Log_Return','Volatility','MA_10','MA_50','RSI','Volume_Change','Open','High','Low']
df = df[cols].dropna().replace([np.inf,-np.inf],np.nan).dropna().reset_index(drop=True)
scaler = StandardScaler()
feat_arr = scaler.fit_transform(df[cols])
target = df['Volatility'].values[1:]
lookback = 21
X_list = [feat_arr[i-lookback:i] for i in range(lookback, len(feat_arr)-1)]
y_list = [target[i] for i in range(lookback, len(feat_arr)-1)]
X = torch.tensor(np.array(X_list),dtype=torch.float32)
y = torch.tensor(np.array(y_list),dtype=torch.float32)
train_size = int(0.8*len(X))
X_train, X_test = X[:train_size], X[train_size:]
y_train, y_test = y[:train_size], y[train_size:]
# -------------------------
# Bayesian LSTM Model Definition
# -------------------------
# Bayesian LSTM: deterministic LSTM + Bayesian final layer
class BayesianLSTM(PyroModule):
def __init__(self, input_dim, hidden_dim, num_layers=1):
super().__init__()
self.lstm = nn.LSTM(input_dim, hidden_dim, num_layers=num_layers, batch_first=True)
self.fc = PyroModule[nn.Linear](hidden_dim, 1)
self.fc.weight = PyroSample(dist.Normal(0., 1.).expand([1, hidden_dim]).to_event(2))
self.fc.bias = PyroSample(dist.Normal(0., 1.).expand([1]).to_event(1))
self.sigma = PyroSample(dist.LogNormal(0., 1.))
def forward(self, x, y=None):
lstm_out, (h_n, _) = self.lstm(x)
h = h_n[-1] # shape: (batch, hidden_dim)
weight = self.fc.weight
bias = self.fc.bias
if weight.dim() == 3: # vectorized case: (S, 1, hidden_dim)
S = weight.size(0)
h_exp = h.unsqueeze(0).expand(S, -1, -1) # (S, batch, hidden_dim)
pred = torch.bmm(h_exp, weight.transpose(1, 2)).squeeze(-1) + bias.unsqueeze(1) # (S, batch)
else:
pred = self.fc(h).squeeze(-1) # (batch)
sigma = self.sigma
if pred.dim() == 2 and sigma.dim() == 1:
sigma = sigma.unsqueeze(1)
with pyro.plate("data", x.shape[0]):
# Always sample "obs", whether or not y is provided.
if y is not None:
obs = pyro.sample("obs", dist.Normal(pred, sigma), obs=y)
else:
obs = pyro.sample("obs", dist.Normal(pred, sigma))
return obs
# -------------------------
# Hyperparameter Tuning with Optuna
# -------------------------
def objective(trial):
pyro.clear_param_store() # clear previous parameters
hidden_dim = trial.suggest_int("hidden_dim", 32, 128)
lr = trial.suggest_loguniform("lr", 1e-3, 1e-1)
model = BayesianLSTM(input_dim=9, hidden_dim=hidden_dim)
guide = AutoNormal(model)
svi = SVI(model, guide, Adam({"lr": lr}), loss=Trace_ELBO())
# Create inner training/validation split from training data (80/20 split)
inner_size = int(0.8 * len(X_train))
X_tr, X_val = X_train[:inner_size], X_train[inner_size:]
y_tr, y_val = y_train[:inner_size], y_train[inner_size:]
train_loader = DataLoader(TensorDataset(X_tr, y_tr), batch_size=64, shuffle=True)
best_val_loss = float("inf")
patience, no_improve = 10, 0
num_epochs = 50
for epoch in range(num_epochs):
for bx, by in train_loader:
svi.step(bx, by)
# Evaluate on validation set
val_loss = 0.
val_loader = DataLoader(TensorDataset(X_val, y_val), batch_size=64)
for bx, by in val_loader:
val_loss += svi.evaluate_loss(bx, by) * bx.size(0)
val_loss /= len(X_val)
if val_loss < best_val_loss:
best_val_loss = val_loss
no_improve = 0
else:
no_improve += 1
if no_improve >= patience:
break
return best_val_loss
study = optuna.create_study(direction="minimize")
study.optimize(objective, n_trials=30)
best_params = study.best_params
print("Best Hyperparameters:", best_params)# -------------------------
# Retrain Final Model with Best Hyperparameters
# -------------------------
pyro.clear_param_store()
final_model = BayesianLSTM(input_dim=9, hidden_dim=best_params["hidden_dim"])
final_guide = AutoNormal(final_model)
final_svi = SVI(final_model, final_guide, Adam({"lr": best_params["lr"]}), loss=Trace_ELBO())
final_epochs = 100
train_loader = DataLoader(TensorDataset(X_train, y_train), batch_size=64, shuffle=True)
for epoch in range(final_epochs):
epoch_loss = sum(final_svi.step(bx, by) for bx, by in train_loader)
if (epoch+1)%10==0:
print(f"Final Model Epoch {epoch+1}: Loss {epoch_loss/len(X_train):.4f}")# -------------------------
# Diagnostics & Forecasting
# -------------------------
adf = adfuller(df['Volatility'])
print(f"ADF Statistic: {adf[0]:.4f}, p-value: {adf[1]:.4f}")
print("Series is likely stationary" if adf[1]<0.05 else "Series is likely non-stationary")Reading the uncertainty
These results come from posterior predictive inference. final_model and final_guide hold the learned distribution over the parameters, and sampling from it a hundred times shows how much the forecast itself would vary. If the number of draws were too small, the uncertainty bands would wobble — but a few hundred samples is enough to trust them.
# Predictive inference
predictive = Predictive(final_model, guide=final_guide, num_samples=100, return_sites=["obs"])
samples = predictive(X_test)["obs"]
y_pred_mean = samples.mean(0).detach().numpy().squeeze()
y_pred_std = samples.std(0).detach().numpy().squeeze()
print(f"Test MSE: {mean_squared_error(y_test.numpy(), y_pred_mean):.4f}")
print(f"Test MAE: {mean_absolute_error(y_test.numpy(), y_pred_mean):.4f}")